import torch.nn as nn
import torch

'''
使用nn.Sequential， 将一系列的层结构打包，形成一个整体
'''


class AlexNet(nn.Module):
    def __init__(self, num_classes=1000, init_weights=False):
        super(AlexNet, self).__init__()
        # 专门用来提取图像特征
        self.features = nn.Sequential(
            nn.Conv2d(3, 48, kernel_size=11, stride=4, padding=2),  # input[3, 224, 224]  output[48, 55, 55]
            nn.ReLU(inplace=True),  # inPlace=True， 增加计算量减少内存使用的一个方法
            nn.MaxPool2d(kernel_size=3, stride=2),                  # output[48, 27, 27]
            nn.Conv2d(48, 128, kernel_size=5, padding=2),           # output[128, 27, 27]
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),                  # output[128, 13, 13]
            nn.Conv2d(128, 192, kernel_size=3, padding=1),          # output[192, 13, 13]
            nn.ReLU(inplace=True),
            nn.Conv2d(192, 192, kernel_size=3, padding=1),          # output[192, 13, 13]
            nn.ReLU(inplace=True),
            nn.Conv2d(192, 128, kernel_size=3, padding=1),          # output[128, 13, 13]
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),                  # output[128, 6, 6]
        )
        # 将全连接层作为一个整体，通过Dropout使其防止过拟合，一般放在全连接层和全连接层之间
        self.classifier = nn.Sequential(
            nn.Dropout(p=0.5),
            nn.Linear(1152, 2048),
            nn.ReLU(inplace=True),
            nn.Dropout(p=0.5),
            nn.Linear(2048, 2048),
            nn.ReLU(inplace=True),
            nn.Linear(2048, num_classes),
        )
        if init_weights:
            self._initialize_weights()

    def forward(self, x):
        print(x.size())
        x = self.features(x)
        x = torch.flatten(x, start_dim=1)  # 展平处理，从channel维度开始展平，（第一个维度为channel）
        x = self.classifier(x)
        return x

    def _initialize_weights(self):
        for m in self.modules():  # 返回一个迭代器，遍历我们网络中的所有模块
            if isinstance(m, nn.Conv2d):  # 判断层结构是否为所给定的层，比如此处判断是否为卷积层
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:  # 此处判断该层偏置是否为空
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):  # 如果是全连接层
                nn.init.normal_(m.weight, 0, 0.01)  # 通过正态分布来给权重赋值
                nn.init.constant_(m.bias, 0)